import sys
import os
import pdfplumber

sys.path.append(os.path.dirname(os.path.abspath(__file__)))

from utils import ArgumentParser, ConfigLoader, LOG
from utils.file_path import FilePathLoader
from model import GLMModel, OpenAIModel
from translator import PDFTranslator
import gradio as gr
from PIL import Image


# Gradio UI
def ui(model_name: str, api_key: str, pdf_file_path: str, file_format: str):
    document_symbol = '\U0001F4C4 选择PDF'  # 📄
    from theme.winter import Winter
    with gr.Blocks(title="PDF Translator", theme=Winter()) as interface:
        gr.Markdown('''# <span style='color:brown'>PDF Translator</span> ''')
        openai_api_key = gr.Textbox(label="APIKey", placeholder="请输入API Key",
                                    value=api_key if api_key else os.getenv("OPENAI_API_KEY"))
        openai_model_name = gr.Textbox(label="模型", placeholder="请输入要使用的OpenAI模型", value=model_name)
        with gr.Row():
            source_pdf = gr.Textbox(label="PDF", placeholder="请选择PDF", value=pdf_file_path)
            source_pdf_button = gr.Button(value=document_symbol)
            source_pdf_button.click(
                get_pdf_file,
                inputs=source_pdf,
                outputs=source_pdf,
                show_progress=False
            )
        with gr.Row():
            res_file_format = gr.Dropdown(choices=["markdown", "pdf"], value=file_format, label="PDF翻译后的文件类型", interactive=True)
            target_language = gr.Dropdown(choices=["简体中文", "英语", "法语", "德语", "意大利语", "韩语"], value="简体中文", label="目标语言", interactive=True)
        translate_button = gr.Button(value="翻译", variant='primary')
        pdf_img = gr.Gallery(label="翻译后的PDF", visible=True)
        res_file_format.change(pdf_img_gallery_visible, inputs=res_file_format, outputs=pdf_img)
        translate_button.click(
            do_translate,
            inputs=[openai_model_name, openai_api_key, source_pdf, res_file_format, target_language],
            outputs=pdf_img
        )

        interface.queue(max_size=8).launch()


# 翻译pdf并将pdf转为图片集进行渲染
def do_translate(model_name: str, api_key: str, pdf_file_path: str, file_format: str, target_language: str):
    # 实例化 PDFTranslator 类，并调用 translate_pdf() 方法
    model = OpenAIModel(model=model_name, api_key=api_key)
    translator = PDFTranslator(model)
    translate_output_file_path = translator.translate_pdf(pdf_file_path, file_format, output_file_path='../tests/translated.pdf', target_language=target_language)
    pil_img_list = []
    if translate_output_file_path and file_format == 'pdf':
        with pdfplumber.open(translate_output_file_path) as pdf:
            for page in pdf.pages:
                img = page.to_image(antialias=True)
                pil_img_list.append(img.original)
    return pil_img_list


# 控制pdf转成图片后的图片集组件可见性
def pdf_img_gallery_visible(res_file_format: str):
    if res_file_format != 'pdf':
        return gr.Gallery.update(visible=False)
    else:
        return gr.Gallery.update(label="翻译后的PDF", object_fit="contain", visible=True)


# 加载pdf路径
def get_pdf_file(file_path):
    path = FilePathLoader()
    return path.get_any_file_path()


if __name__ == "__main__":
    argument_parser = ArgumentParser()
    args = argument_parser.parse_arguments()
    # config_loader = ConfigLoader(args.config)
    config_loader = ConfigLoader("..\config.yaml")

    config = config_loader.load_config()

    model_name = args.openai_model if args.openai_model else config['OpenAIModel']['model']
    api_key = args.openai_api_key if args.openai_api_key else config['OpenAIModel']['api_key']

    pdf_file_path = args.book if args.book else config['common']['book']
    file_format = args.file_format if args.file_format else config['common']['file_format']

    ui(model_name, api_key, pdf_file_path, file_format)
